import jieba
jieba.add_word('林可',500)
from rank_bm25 import BM25Okapi
from AiAssistant._common import calculate_execution_time
#pip install rank_bm25  -i https://pypi.tuna.tsinghua.edu.cn/simple

@calculate_execution_time
def tokenization(corpus):
    tokenized_corpus = [jieba.lcut(doc) for doc in corpus]
    return tokenized_corpus
"""
tokenization 的当前/Max执行时间为:2.5168 秒/2.5168秒。Total执行时间为: 2.517 秒,平均 2.517
总知识数 29674
init_bm25 的当前/Max执行时间为:0.1498 秒/0.1498秒。Total执行时间为: 0.15 秒,平均 0.15
build_bm_model 的当前/Max执行时间为:2.6697 秒/2.6697秒。Total执行时间为: 2.67 秒,平均 2.67

bm25_search 的当前/Max执行时间为:0.0269 秒/0.077秒。Total执行时间为: 0.77 秒,平均 0.02851851851851852
"""
@calculate_execution_time
def init_bm25(tokenized_corpus):
    bm25 = BM25Okapi(tokenized_corpus)
    return bm25

@calculate_execution_time
def build_bm_model(corpus):
    tokenized_corpus = tokenization(corpus)
    print('总知识数',len(tokenized_corpus))
    bm25 = init_bm25(tokenized_corpus)
    return bm25

@calculate_execution_time
def build_char_bm_model(corpus):
    tokenized_corpus = [list(u) for u in corpus]
    print('总知识数',len(tokenized_corpus))
    bm25 = init_bm25(tokenized_corpus)
    return bm25


@calculate_execution_time
def bm25_search(bm25,corpus,query,sw,topk=100,char=False):
    tnp = """
1-1
1-2
1-6
10-3
10-5
10-6
11-2
11-3
12-1
12-2
13-2
15-1
16-3
2-1
2-2
2-4
2-5
2-6
3-2
3-4
3-6
3-7
5-4
5-6
6-1
6-2
6-3
6-4
6-5
6-9
7-1
7-2
8-1
8-2
8-3
8-4
9-2
9-3
9-6
15-2
6-7
3-5
2-8
5-7
2-3
1-5
7-4
16-1
1-3
13-3
1-4
10-4
    """
    entity = tnp.splitlines()
    entity = [e for e in entity if '-' in e]
    if char:
        tokenized_query2 = list(query)
    else:
        tokenized_query2 = jieba.lcut(query)
    tokenized_query = [e for e in tokenized_query2 if e not in sw]
    if len(tokenized_query) !=len(tokenized_query2):
        print('减少')

    # 计算查询与每个文档的BM25分数
    doc_scores = bm25.get_scores(tokenized_query)

    # 将文档和分数配对，然后按分数排序
    scored_corpus = sorted(zip(corpus, doc_scores), key=lambda x: x[1], reverse=True)
    if any(e in query for e in entity):
        en_query = [e for e in entity if e in query]
        print('entity',en_query)
        return [d for d, s in scored_corpus if any(e in d for e in en_query)][:topk]


    return [d for d,s in scored_corpus ][:topk]


class RankerBM25:
    def __init__(self,corpus,char= False):
        if char:
            self.model = build_char_bm_model(corpus)
        else:
            self.model = build_bm_model(corpus)
        self.corpus = corpus

    def sort(self,query,sw=[],topk=100,char = False):
        return bm25_search( self.model, self.corpus, query, sw, topk=topk,char=char)
